# whisper_asr.py
from typing import Optional
import torch
import whisper
from opencc import OpenCC

from .asr_base import BaseASR
from .asr_registry import register_asr
import gc


@register_asr("whisper")
class WhisperASR(BaseASR):
    def __init__(self):
        super().__init__()
        self.model_name = "small"
        self.device = "cpu"
        self.cc = OpenCC('t2s')  # 简繁转换器
        self._model = None

    @property
    def model(self):
        """懒加载模型"""
        if self._model is None:
            # 加载新模型
            print(f"[WhisperASR] 正在加载模型 '{self.model_name}'，设备: {self.device}")
            self._model = whisper.load_model(self.model_name).to(self.device)

            print(f"[WhisperASR] 模型 '{self.model_name}' 加载完成")
            return self._model

    def transcribe(self, audio_data=None, audio_path: Optional[str] = None) -> str:
        if audio_data is None and audio_path is None:
            raise ValueError("audio_data 或 audio_path 必须提供一个")

        print("[WhisperASR] 开始转录...")

        # 设置 fp16（仅 GPU 支持）
        # fp16 = self.device == "cuda"

        result = self.model.transcribe( # type: ignore
                audio_path if audio_path else audio_data.copy(),
                language="zh",
                fp16=False,
                word_timestamps=False
            )

        # 简繁转换
        simplified_text = self.cc.convert(result["text"])
        return simplified_text

    def cleanup(self):
        """主动释放模型资源"""
        if self._model is not None:
            print(f"[WhisperASR] 正在释放模型 '{self.model_name}' 资源")
            del self._model
            self._model = None
            print("[WhisperASR] 模型已从内存中删除")

            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                print("[WhisperASR] CUDA 缓存已清空")

            gc.collect()
            print("[WhisperASR] 资源释放完成")

    def __del__(self):
        """析构函数兜底"""
        if self._model is not None:
            print(f"[WhisperASR] 析构函数调用，释放模型 '{self.model_name}' 资源")
            self.cleanup()